package com.chatplus.application.aiprocessor.util;

import cn.hutool.extra.spring.SpringUtil;
import com.chatplus.application.common.constant.RedisPrefix;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.common.util.IdGenerator;
import com.chatplus.application.common.util.MessageTokenUtil;
import com.chatplus.application.common.util.spring.SpringCtxUtils;
import com.chatplus.application.aiprocessor.provider.ChatAiProcessorServiceProvider;
import com.chatplus.application.aiprocessor.channel.chat.ChatAiProcessorService;
import com.chatplus.application.aiprocessor.handler.dto.ChatRecordMessage;
import com.chatplus.application.enumeration.AiPlatformEnum;
import org.redisson.api.RBucket;
import org.redisson.api.RedissonClient;
import org.springframework.web.socket.WebSocketSession;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;

/**
 * 会话工具类，比如发送消息。保存历史记录等等
 * 基于Spring WebSocket
 */
public class ChatAiUtilHandler {
    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(ChatAiUtilHandler.class);
    protected final WebSocketSession session;
    /**
     * 平台
     */
    protected final AiPlatformEnum platform;
    /**
     * 用户发送的问题
     */
    protected final String prompt;

    public ChatAiUtilHandler(WebSocketSession session, AiPlatformEnum platform, String prompt) {
        this.session = session;
        this.platform = platform;
        this.prompt = prompt;
    }


    public void saveChatHistory(String reply, Long token) {
        String modelName = session.getAttributes().get("modelName").toString();
        Long promptTokens = MessageTokenUtil.countMessageTokens(prompt, modelName);
        Long replyTokens = token == null ? MessageTokenUtil.countMessageTokens(reply, modelName) : token;
        ChatAiProcessorServiceProvider chatAiProcessorServiceProvider = SpringCtxUtils.getBean(ChatAiProcessorServiceProvider.class);
        ChatAiProcessorService aiProcessor = chatAiProcessorServiceProvider.getAiProcessorService(platform);
        // 保存记录
        aiProcessor.saveChatItemRecord(prompt);
        //参数1为终端ID
        //参数2为数据中心ID
        Long itemId = IdGenerator.generateLongId();
        aiProcessor.saveChatHistory(itemId, "prompt", promptTokens, prompt);
        aiProcessor.saveChatHistory(itemId, "reply", replyTokens, reply);
        // 保存缓存
        RedissonClient redissonClient = SpringUtil.getBean(RedissonClient.class);
        String chatId = session.getAttributes().get("chat_id").toString();
        RBucket<List<ChatRecordMessage>> rBucket = redissonClient.getBucket(RedisPrefix.TURN_RIGHT_REDIS_PREFIX + chatId);
        List<ChatRecordMessage> recordMessageList = rBucket.isExists() ? rBucket.get() : new ArrayList<>();
        ChatRecordMessage chatRecordMessage = new ChatRecordMessage();
        chatRecordMessage.setPrompt(prompt);
        chatRecordMessage.setReply(reply);
        recordMessageList.add(chatRecordMessage);
        rBucket.set(recordMessageList, Duration.ofMinutes(30));
    }

    public void sendMessage(String message) {
        ChatAiProcessorServiceProvider chatAiProcessorServiceProvider = SpringCtxUtils.getBean(ChatAiProcessorServiceProvider.class);
        ChatAiProcessorService aiProcessor = chatAiProcessorServiceProvider.getAiProcessorService(platform);
        aiProcessor.sendMessage(message);
    }

    /**
     * 回复错误消息
     *
     * @param message   消息
     * @param haveStart 是否包含开始开始标志，包不包含，得看StreamEventSourceListener.onOpen有没有连接成功
     *                  ，如果是在onEvent里面发生错误的，一般是传false的，除非是需要换输出行的
     */
    public void replyErrorMessage(String message, boolean haveStart) {
        ChatAiProcessorServiceProvider chatAiProcessorServiceProvider = SpringCtxUtils.getBean(ChatAiProcessorServiceProvider.class);
        ChatAiProcessorService aiProcessor = chatAiProcessorServiceProvider.getAiProcessorService(platform);
        if (haveStart) {
            aiProcessor.replyMessage(message);
        } else {
            aiProcessor.replyNotStartMessage(message);
        }

    }
}
